#!/usr/bin/perl -w
use strict;
use Getopt::Long;
use lib "/net/cpp-group/Leo/bin";
use parse_fasta;
#88888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Usage

#88888888888888888888888888888888888888888888888888888888888888888888888888888888
my $usage = <<"USAGE";

USAGE:

    filter_gaps_from_aligned_fasta.pl
					--prot_in		prot.fa
					--cdna_in		cdna.fa
                    --prot_out		prot_new.fa
					--cdna_out		cdna_new.fa
                    [--gaps_percent PERCENT
                     --gaps_max_seq CNT_SEQ
					--regex "(\\S+)"			]

        This script accepts sequences in the FASTA format and remove all residues
        and codons where there are alignment gaps in
        over PERCENT% or CNT_SEQ of the peptides sequences.
		PERCENT defaults to 90
		CNT_SEQ defaults to 2

USAGE

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Get options

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
# mandatory




# optional parameters
my $help = undef;
my $gaps_percent = 90;
my $gaps_max_seq = 2;
my $prot_file = undef;
my $cdna_file = undef;
my $prot_out_file = undef;
my $cdna_out_file = undef;
my $regex = '(\S+)';
GetOptions('help' => \$help,
		   'gaps_percent=s' => \$gaps_percent,
		   'gaps_max_seq=s' => \$gaps_max_seq,
		   'prot_in=s' => \$prot_file,
		   'cdna_in=s' => \$cdna_file,
		   'prot_out=s' => \$prot_out_file,
		   'cdna_out=s' => \$cdna_out_file,
		   'regex=s' => \$regex,
		   );

die $usage if ($help);
die $usage unless ($prot_file && $cdna_file && $cdna_out_file && $prot_out_file);


#_________________________________________________________________________________________
#
#		Read peptides
#
sub read_peptides(\@\@$)
{
	my ($accessions, $prot, $prot_file) = @_;
	sub save_prot_callback($$$$$$)
	{
		my ($acc, $seq_lines, $data) = @_[0, 1, 5];
		my ($prot, $accessions) = @$data;
		$acc = $1 if ($acc =~ /$regex/);

		# save sequences and accessions, the former as an array of letters
		my $sequence = join "", @$seq_lines;
		push(@$prot, [unpack("C*", $sequence)]);
		push(@$accessions, $acc);
	}

	# open file and pass to callback
	open SEQ, $prot_file or die "Error:\n\tCould not open $prot_file\n$!";
	my $cnt_seq =parse_fasta::parse_sequences(  *SEQ, \&save_prot_callback,
												[$prot, $accessions]);
	print STDERR "\t$cnt_seq\tpeptides parsed.\n";
}
#
#
#_________________________________________________________________________________________


#_________________________________________________________________________________________
#
#		Read cDNA
#
sub read_cdna(\@\@$)
{
	my ($accessions, $cdna, $cdna_file) = @_;

	# turn accessions into a hash for quick lookup
	my %accessions;
	@accessions{@$accessions}=(0..$#$accessions);

	sub save_cdna_callback($$$$$$)
	{
		my ($acc, $seq_lines, $data) = @_[0, 1, 5];
		$acc = $1 if ($acc =~ /$regex/);
		
		my ($accessions, $cdna) = @$data;

		# make sure corresponding peptide parsed
		return unless exists $accessions->{$acc};
		my $seq = join "", @$seq_lines;
		$cdna->[$accessions->{$acc}] = [unpack("A3" x (length($seq) / 3) , $seq)];
	}

	# open file and pass to callback
	open SEQ, $cdna_file or die "Error:\n\tCould not open $cdna_file\n$!";
	my $cnt_seq =parse_fasta::parse_sequences(  *SEQ, \&save_cdna_callback, [\%accessions, $cdna]);

	print STDERR "\t$cnt_seq\tcdna sequences parsed.\n";
}
#
#
#_________________________________________________________________________________________


#_________________________________________________________________________________________
#
#		Check sequences length
#

sub check_len($$$$)
{
	my ($accession, $common_len, $len, $name) = @_;
	if ($common_len != $len)
	{
		die "Error:\n\t".
				"The first peptide sequence has $common_len residues ".
				"but the $name sequence $accession encodes $len residues.\n".
				"All aligned sequences should have the same length!\n";

	}
}
sub check_seq_length(\@\@\@)
{
	my ($prot, $cdna, $accessions) = @_;

	die "Error:\n\tNo peptide sequences supplied\n" unless  @$prot;
	#
	#	check same length
	#
	my $len = @{$prot->[0]};
	for (my $i = 0; $i < @$prot; ++$i)
	{
		check_len($accessions->[$i], $len, scalar @{$prot->[$i]}, "peptide");
	}
	for (my $i = 0; $i < @$cdna; ++$i)
	{
		check_len($accessions->[$i], $len, scalar @{$cdna->[$i]}, "cdna");
	}
	return $len;

}
#
#
#_________________________________________________________________________________________

#_________________________________________________________________________________________
#
#		copy_ungapped
#
#			copy proteins and cdna which have fewer than the specified number of gaps
#			at each position

sub copy_ungapped(\@\@\@\@$$)
{
	my ($prot, $cdna, $new_prot, $new_cdna, $len, $max_gaps) = @_;
	my $cnt_removed_gaps = 0;
	my $gap = ord('-');
	for (my $pos = 0; $pos < $len; ++$pos)
	{
		# count gaps at position $pos
		my $cnt_gaps = 0;
		for my $seq (@$prot)
		{
			++$cnt_gaps if ($seq->[$pos] == $gap);
		}

		# skip if too many gaps
		if ($cnt_gaps >= $max_gaps)
		{
			$cnt_removed_gaps++;
			next;
		}

		# add residue and codon at this position to new sequences
		for (my $i = 0; $i < @$prot; ++$i)
		{
			$new_prot->[$i].= chr($prot->[$i][$pos]);
			$new_cdna->[$i].= $cdna->[$i][$pos];
		}
	}
	return $cnt_removed_gaps;
}
#
#
#_________________________________________________________________________________________


#_________________________________________________________________________________________
#
#		print_sequences
#
sub print_sequences(\@\@$)
{
	my ($sequences, $accessions, $file_name) = @_;
	open FILE, ">$file_name" or die "Error:\n\tCould not open $file_name\n$!";
	for (my $i = 0; $i < @$sequences; ++$i)
	{
		print FILE ">$accessions->[$i]\n$sequences->[$i]\n";
	}


}
#
#
#_________________________________________________________________________________________




	#	Read peptides and cdna
	my @accessions;
	my @prot;
	read_peptides(@accessions, @prot, $prot_file);
	my @cdna;
	read_cdna(@accessions, @cdna, $cdna_file);
	die "Error:\n\tThe cdna sequences for some peptides were not found\n"
				unless scalar @cdna == scalar @prot;

	#	Check lengths
	my $len = check_seq_length(@prot, @cdna, @accessions);
	
	
	# 	Count number of gaps to disqualify position
	my $seq_count = @prot;
	my $max_gaps = $seq_count - int($seq_count * $gaps_percent / 100) + 1;
	$max_gaps = $gaps_max_seq if ($max_gaps > $gaps_max_seq);
	print STDERR "\tGaps in $max_gaps or more prot will be discarded\n";




	# 	Prepare arrays for new sequences
	my @new_prot;
	@new_prot = ('') x scalar @prot;
	my @new_cdna;
	@new_cdna = ('') x scalar @prot;

	my $cnt_removed_gaps = copy_ungapped(@prot, @cdna, @new_prot, @new_cdna, $len, $max_gaps);
	print STDERR "\t$cnt_removed_gaps\tgap positions removed.\n";

	print_sequences(@new_prot, @accessions, $prot_out_file);
	print_sequences(@new_cdna, @accessions, $cdna_out_file);
